import math
import maya.cmds as cmds
import AntsCrowd.transform as transform

mode = 1
walkFreq = 3
upSpeed = 20
radL = 0
radR = 0
maxTime = 60
L_offSet = 0
R_offset = 0

try :    
    cmds.delete( dellist )
    dellist = []
except :
    dellist = []

def distance(a,b):    
    tot = 0
    for i in range(3):
        tot +=(a[i]-b[i])**2
    return math.sqrt(tot)

def getHich( tg ):    
    main = [ m for m in cmds.listRelatives( tg,c = 1) if 'main' in m]    
    leg_grp = [ l for l in cmds.listRelatives( tg,c = 1) if 'leg_grp' in l]
    null_grp = [ l for l in cmds.listRelatives( main ,c = 1) if 'null' in l]
    legs = cmds.listRelatives( leg_grp[0],c = 1)
    return {'main':main , 'legs':legs , 'null':null_grp , 'res':1}
    

#############################################################################

#############################################################################



def setLiftPos( null , ctrl , liftValue  ):
    
    mat = cmds.getAttr( null + '.worldMatrix' )
    cmds.xform( ctrl , m = mat )
    nullY = transform.getMtoT( cmds.getAttr( null+'.worldMatrix'  ) )[1]
    cmds.setAttr( ctrl +'.ty' , liftValue * upSpeed +  nullY   )     
    cmds.setKeyframe( ctrl +'.t')
    cmds.setKeyframe( ctrl +'.r' )
    
    
    
def setStickyPos( ctrl , mat ):
    
    cmds.xform( ctrl , m = mat )
    cmds.setKeyframe(ctrl +'.t' )
    cmds.setKeyframe(ctrl +'.r' )

def antwalk( ant , walkFreq = 3 , upSpeed = 10 , maxTime = 60 ):
    
    global radL
    global radR 
    
    upSpeed = walkFreq*40
    if ant['res'] != 1 : return    
    
    ant_main = ant['main'][0]
    ant_lfft_ctrl = [a for a in ant['legs'] if 'lfft' in a ][0]
    ant_lfmd_ctrl = [a for a in ant['legs'] if 'lfmd' in a ][0]
    ant_lfbk_ctrl = [a for a in ant['legs'] if 'lfbk' in a ][0]
    ant_rtft_ctrl = [a for a in ant['legs'] if 'rtft' in a ][0]
    ant_rtmd_ctrl = [a for a in ant['legs'] if 'rtmd' in a ][0]
    ant_rtbk_ctrl = [a for a in ant['legs'] if 'rtbk' in a ][0]
    
    ant_lfft_null = [a for a in ant['null'] if 'lfft' in a ][0]
    ant_lfmd_null = [a for a in ant['null'] if 'lfmd' in a ][0]
    ant_lfbk_null = [a for a in ant['null'] if 'lfbk' in a ][0]
    ant_rtft_null = [a for a in ant['null'] if 'rtft' in a ][0]
    ant_rtmd_null = [a for a in ant['null'] if 'rtmd' in a ][0]
    ant_rtbk_null = [a for a in ant['null'] if 'rtbk' in a ][0]
       
    for x in range( maxTime ):
        cmds.currentTime( x )
        if x == 0 :            
            for x in [ 'lf' , 'rt' ]:
                for y in [ 'ft' , 'md' , 'bk' ]:
                    mat = cmds.getAttr( eval('ant_'+x+y+'_null') + '.worldMatrix' ) 
                    cmds.xform( eval('ant_'+x+y+'_ctrl') , m = mat )
                    cmds.setKeyframe(eval('ant_'+x+y+'_ctrl') +'.t' )
                    cmds.setKeyframe(eval('ant_'+x+y+'_ctrl') +'.r' )
            L_offSet = 0                                  
            R_offset = 0

        else :
            if max( 0 , math.sin( radL ) ) != 0  :
                setLiftPos(  ant_lfft_null , ant_lfft_ctrl , math.sin( radL )  )
                setLiftPos(  ant_rtmd_null , ant_rtmd_ctrl , math.sin( radL ) )
                setLiftPos(  ant_lfbk_null , ant_lfbk_ctrl , math.sin( radL ) )
               
                if L_offSet == 0 :
                    lfmd_mat = cmds.getAttr( ant_lfmd_null + '.worldMatrix' )
                    rtft_mat = cmds.getAttr( ant_rtft_null + '.worldMatrix' )
                    rtbk_mat = cmds.getAttr( ant_rtbk_null + '.worldMatrix' )
                    
                setStickyPos( ant_lfmd_ctrl , lfmd_mat )
                setStickyPos( ant_rtft_ctrl , rtft_mat )
                setStickyPos( ant_rtbk_ctrl , rtbk_mat ) 
                L_offSet = 1            
                R_offset = 0
                
            else :
                
                setLiftPos(  ant_rtft_null , ant_rtft_ctrl ,
                           max( 0 , math.sin( radR ) )   )                
                setLiftPos(  ant_lfmd_null , ant_lfmd_ctrl ,
                           max( 0 , math.sin( radR ) )   )
                setLiftPos(  ant_rtbk_null , ant_rtbk_ctrl ,
                           max( 0 , math.sin( radR ) )   )

                radR += math.pi / walkFreq
                
                if R_offset == 0 :
                    lfft_mat = cmds.getAttr( ant_lfft_null + '.worldMatrix' )
                    lfbk_mat = cmds.getAttr( ant_lfbk_null + '.worldMatrix' )
                    rtmd_mat = cmds.getAttr( ant_rtmd_null + '.worldMatrix' )
                setStickyPos( ant_lfft_ctrl , lfft_mat )
                setStickyPos( ant_lfbk_ctrl , lfbk_mat )
                setStickyPos( ant_rtmd_ctrl , rtmd_mat ) 
                L_offSet = 0           
                R_offset = 1
    
            radL += math.pi / walkFreq
